'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: dataset_handler.py
* @Time: 2025/9/25
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass
from uuid import UUID

from flask import request
from flask_login import login_required, current_user
from injector import inject

from internal.core.file_extractor import FileExtractor
from internal.model import UploadFile
from internal.schema.dataset_schema import CreateDataReq, GetDatasetResp, GetDatasetWithPageReq, GetDatasetWithPageResp, \
    HitReq, GetDatasetQueriesResp
from internal.service import DatasetService, EmbeddingsService, JiebaService
from pkg.paginator.paginator import PageModel
from pkg.response import validate_error_json, success_message, success_json
from pkg.sqlachemy.sqlalchemy import SQLAlchemy


@inject
@dataclass
class DatasetHandler:
    """知识库http请求处理器"""
    db:SQLAlchemy
    data_service: DatasetService
    embeddings_service:EmbeddingsService
    jieba_service:JiebaService
    file_extractor:FileExtractor


    @login_required
    def create_dataset(self):
        """创建知识库http请求"""
        #1、提取请求校验
        req = CreateDataReq()
        if not req.validate():
            return validate_error_json(req.errors)
        #2、调用服务创建知识库
        self.data_service.create_dataset(req, current_user)

        #3、返回成功调用提示
        return success_message("创建知识库成功！")


    @login_required
    def get_dataset(self, dataset_id:UUID):
        """根据传递的知识库id 获取知识库详情"""
        dataset = self.data_service.get_dataset(dataset_id, current_user)
        resp = GetDatasetResp()
        return success_json(resp.dump(dataset))


    @login_required
    def update_dataset(self, dataset_id: UUID):
        # 1、提取请求校验
        req = CreateDataReq()
        if not req.validate():
            return validate_error_json(req.errors)
        #调用服务并更新知识库
        self.data_service.update_dataset(dataset_id,req, current_user)

        return success_message("更新知识库成功！")


    @login_required
    def get_datasets_with_page(self):
        """获取知识库分页+搜索数据"""
        req = GetDatasetWithPageReq(request.args)
        if not req.validate():
            return validate_error_json(req.errors)

        #2、调用服务获取分页数据
        datasets, paginator = self.data_service.get_dataset_with_page(req, current_user)

        resp = GetDatasetWithPageResp(many=True)
        return success_json(PageModel(list=resp.dump(datasets), paginator=paginator))


    @login_required
    def embeddings_query(self):
        """向量查询测试"""
        upload_file = self.db.session.query(UploadFile).get("cee3a4a8-b1eb-47d6-8ba8-ce0427168c34")
        content = self.file_extractor.load(upload_file,True)
        return success_json({"vector":content})


    @login_required
    def hit(self, dataset_id: UUID):
        """根据传递的知识库id + 检索参数执行召回测试"""
        req = HitReq()
        if not req.validate():
            return validate_error_json(req.errors)

        #2、调用服务执行检索策略
        hit_result = self.data_service.hit(dataset_id, req, current_user)
        return success_json(hit_result)


    @login_required
    def get_data_queries(self, dataset_id:UUID):
        """根据传递的知识库id获取最近的10条查询记录"""
        dataset_queries = self.data_service.get_dataset_queries(dataset_id, current_user)
        resp = GetDatasetQueriesResp(many=True)
        return success_json(resp.dump(dataset_queries))


    @login_required
    def delete_dataset(self, dataset_id: UUID):
        """根据传递的知识库id删除知识库"""
        self.data_service.delete_dataset(dataset_id, current_user)
        return success_message("删除知识库成功")
